package org.example;

import org.example.util.BiMap;
import org.example.util.IndexUtil;
import org.xbill.DNS.*;
import org.xbill.DNS.Record;

import java.net.*;
import java.util.*;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

public class SimpleDnsServer implements Runnable {

    private static final int DNS_PORT = 53;
    private final InetAddress bindAddress;
    private final InetAddress toAddress;

    private final Resolver defaultResolver;

    public static final BiMap<String,InetAddress> biMap = new BiMap<>();
    private final List<String> allowedTlds = Arrays.asList("docker.com");

    private ExecutorService executor = Executors.newVirtualThreadPerTaskExecutor();


    public SimpleDnsServer(String bindAddress,String toAddress) {
        try {
            this.bindAddress = InetAddress.getByName(bindAddress);
            this.toAddress = InetAddress.getByName(toAddress);
            this.defaultResolver = new SimpleResolver("8.8.8.8");
            this.defaultResolver.setTCP(false);
            this.defaultResolver.setPort(53);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    private InetAddress incrementIp(InetAddress base, int increment) throws UnknownHostException {
        byte[] addr = base.getAddress();
        for (int i = addr.length - 1; i >= 0 && increment > 0; i--) {
            int val = (addr[i] & 0xFF) + increment;
            addr[i] = (byte) (val & 0xFF);
            increment = val >> 8;
        }
        return InetAddress.getByAddress(addr);
    }

    private boolean matchesAllowedTld(String domain) {
        // 域名去掉末尾的点
        domain = domain.toLowerCase(Locale.ROOT);
        if (domain.endsWith(".")) {
            domain = domain.substring(0, domain.length() - 1);
        }

        for (String tld : allowedTlds) {
            if (domain.endsWith(tld)) {
                return true;
            }
        }
        return false;
    }

    @Override
    public void run() {
        try (DatagramSocket socket = new DatagramSocket(new InetSocketAddress(bindAddress, DNS_PORT))) {
            byte[] buffer = new byte[512];
            System.out.println("✅ DNS 服务器已启动，监听 " + bindAddress.getHostAddress() + ":" + DNS_PORT);

            while (true) {
                DatagramPacket requestPacket = new DatagramPacket(buffer, buffer.length);
                socket.receive(requestPacket);

                // 克隆数据包内容，避免线程间数据竞争
                byte[] requestData = Arrays.copyOf(requestPacket.getData(), requestPacket.getLength());
                InetAddress clientAddress = requestPacket.getAddress();
                int clientPort = requestPacket.getPort();

                executor.submit(() -> {
                    try {
                        Message query = new Message(requestData);
                        Record question = query.getQuestion();
                        Name name = question.getName();
                        int type = question.getType();

                        String domain = name.toString(true).toLowerCase();
                        System.out.println("📥 查询域名: " + domain + " | 类型: " + Type.string(type));

                        Message response;

                        if (type == Type.A && matchesAllowedTld(domain)) {
                            InetAddress mappedAddress = biMap.getValue(domain);
                            if (mappedAddress == null) {
                                //双重校验防止重复获取
                                synchronized (domain.intern()){
                                    mappedAddress = biMap.getValue(domain);
                                    if(mappedAddress == null){
                                        int availableIndex = IndexUtil.getAvailableIndex();
                                        mappedAddress = incrementIp(toAddress, availableIndex);
                                        biMap.put(domain, mappedAddress);
                                        System.out.println("🔗 为域名分配 IP: " + domain + " → " + mappedAddress.getHostAddress());
                                    }
                                }
                            }

                            response = new Message(query.getHeader().getID());
                            response.getHeader().setFlag(Flags.QR);
                            response.addRecord(question, Section.QUESTION);

                            ARecord aRecord = new ARecord(name, DClass.IN, 60, mappedAddress);
                            response.addRecord(aRecord, Section.ANSWER);
                        } else {
                            response = defaultResolver.send(query);
                        }

                        byte[] responseData = response.toWire();
                        DatagramPacket responsePacket = new DatagramPacket(
                                responseData, responseData.length,
                                clientAddress, clientPort
                        );

                        socket.send(responsePacket);
                    } catch (Exception e) {
                        e.printStackTrace();
                        System.err.println("❌ 处理查询失败: " + e.getMessage());
                    }
                });
            }
        } catch (Exception ex) {
            ex.printStackTrace();
        }
    }
}
